# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import logging
from typing import Callable

from coin.exchange.base.order_gateway import Order
from coin.base.timestamp import get_timestamp


class OrderManagerError(ValueError):
  pass


class OrderManager(object):
  def __init__(self):
    self._orders = {}  # all the added orders will be put in this dict.
    self._working_orders = {}
    self._orders_by_exchange_order_id = {}
    self._logger = logging.getLogger(__name__)

  def __contains__(self, order):
    return order.order_id in self._orders

  def _check_addable(self, new_order):
    if new_order.order_id is None:
      raise OrderManagerError('order id is None')
    if new_order.order_id in self._orders:
      raise OrderManagerError('Duplicated order id')
    if new_order.internal.exchange_order_id in self._orders_by_exchange_order_id:
      raise OrderManagerError('Duplicated exchange_order_id')

  def add_order(self, order: Order):
    logging.info('add order:%s', order.order_id)
    self._check_addable(order)
    self._orders[order.order_id] = order
    if not (order.internal.cancel_confirmed or order.internal.fully_filled):
      self._working_orders[order.order_id] = order
    if order.internal.exchange_order_id is not None:
      self._orders_by_exchange_order_id[order.internal.exchange_order_id] = order

  def remove_order(self, order: Order, *, ignore_error=False):
    logging.info('remove_order:%s', order.order_id)
    if order.order_id not in self._orders:
      raise OrderManagerError('remove_order error: order not found! %s' % order)

    if ignore_error:
      self._orders.pop(order.order_id, None)
    else:
      self._orders.pop(order.order_id)

    self._working_orders.pop(order.order_id, None)
    self._orders_by_exchange_order_id.pop(order.internal.exchange_order_id, None)

  def is_dead_order(self, order):
    assert order.order_id in self._orders, 'Unknown order, %s' % order
    return order.internal.fully_filled or order.internal.cancel_confirmed

  def _set_exchange_order_id(self, order, exchange_order_id):
    # This method can only be called once for each order.
    assert isinstance(exchange_order_id, str), exchange_order_id
    assert order.accepted is False, order
    assert order.internal.exchange_order_id is None, order
    assert exchange_order_id not in self._orders_by_exchange_order_id, (
        exchange_order_id, order, self._orders_by_exchange_order_id[exchange_order_id])

    order.accepted = True
    order.internal.exchange_order_id = exchange_order_id
    self._orders_by_exchange_order_id[exchange_order_id] = order

  def _set_cancel_confirmed(self, order, cancel_confirmed):
    assert cancel_confirmed in (True, False)
    assert order.internal.exchange_order_id is not None, order
    order.internal.cancel_confirmed = cancel_confirmed
    self._working_orders.pop(order.order_id, None)  # May removed already.

  def _set_fully_filled(self, order, fully_filled):
    assert fully_filled in (True, False)
    assert order.internal.exchange_order_id is not None, order
    order.internal.fully_filled = fully_filled
    self._working_orders.pop(order.order_id, None)  # May removed already.

  def set_order(
      self,
      order: Order,
      *,
      exchange_order_id=None,
      fully_filled=None,
      cancel_confirmed=None,
  ):
    assert order.order_id is not None and order.order_id in self._orders, order

    # Only one of the parameters is specified, all others are None.
    params_set = {exchange_order_id, fully_filled, cancel_confirmed}
    assert None in params_set and len(params_set) == 2

    if exchange_order_id is not None:
      self._set_exchange_order_id(order, exchange_order_id)
    elif fully_filled is not None:
      self._set_fully_filled(order, fully_filled)
    elif cancel_confirmed is not None:
      self._set_cancel_confirmed(order, cancel_confirmed)

  def get_working_order(self):
    return list(self._working_orders.values())

  def filter(self, predicate: Callable, transform=None):
    transform = transform or (lambda x: x)
    res = []
    for order in self._orders.values():
      if predicate(order):
        res.append(transform(order))
    return res

  def get(self, *, order_id=None, exchange_order_id=None):
    # python style xor
    assert (order_id is None) != (exchange_order_id is None)

    if order_id is not None:
      return self._working_orders.get(order_id, None)
    elif exchange_order_id is not None:
      return self._orders_by_exchange_order_id.get(exchange_order_id, None)
    else:
      raise OrderManagerError('Invalid parameters for get method. %s, %s' %
                              (order_id, exchange_order_id))

  def copy_from_exchange_open_orders(
      self,
      exchange_order_dict,
      products,
      *,
      copy_exclude_time_threshold=5e9,
      max_num_not_seen_times=3,
  ):
    # local working order list and exchange working order list may not be the same,
    # but we assume that except for those newly submit-accepted orders
    # and those newly cancel-confirmed orders, all other orders should be the same.

    timestamp = get_timestamp()
    for order in self.get_working_order().copy():
      if order.product not in products:
        continue

      if order.internal.exchange_order_id not in exchange_order_dict:
        diff_time = timestamp - order.internal.timestamp
        if diff_time > copy_exclude_time_threshold:
          order.internal.num_not_seen_times += 1
          if order.internal.num_not_seen_times > max_num_not_seen_times:
            self._logger.error('Order removed. %s (diff_time = %s)', order.order_id, diff_time)
            self.remove_order(order)

    for exchange_order_id, order in exchange_order_dict.items():
      prev_order = self.get(exchange_order_id=exchange_order_id)
      if prev_order is None:
        prev_order = self.get(order_id=order.order_id)
        if prev_order is None:
          self.add_order(order)
        else:
          self.remove_order(prev_order)
          order.order_id = prev_order.order_id
          self.add_order(order)
      else:
        if self.is_dead_order(prev_order):
          self._logger.warning('dead order (%s, %s) appeared in exchange working order list!',
                               order.order_id,
                               prev_order.order_id)
        else:
          if prev_order.cancel_sent and timestamp - prev_order.internal.cancel_sent_time > copy_exclude_time_threshold:
            self._logger.warning('Previous order cancel (%s) failed!', order.order_id)
            prev_order.cancel_sent = False
